Greptile Summary

This PR adapts TransformerEngine to leverage cuDNN's recent enhancement that allows returning any subset of {Stats, SumExp, Max}. The implementation now always retrieves Stats from cuDNN (which equals log(SumExp)+Max) and additionally obtains the Max tensor when return_max_logit=True.

Key changes:

• Set generate_stats=true to force cuDNN to always return the Stats tensor needed for backward pass
• Removed all references to SumExp tensor since it's no longer needed (Stats is returned directly)
• Eliminated manual computation of Stats = log(SumExp) + Max in Python layer
• Renamed generate_max_sum_exp to return_max_logit in FADescriptor_v1 for clarity
• Updated documentation across all API functions to reflect new semantics
• Adjusted tensor ordering in output: now returns Stats first, then Max (when requested)

All previous review feedback has been addressed, including updated comments in fused_attn.py and documentation in fused_attn.h.

Confidence Score: 5/5

• This PR is safe to merge with minimal risk - changes are well-structured and all previous feedback has been addressed
• The implementation correctly adapts to cuDNN's new capability, maintains consistency across all layers (C++, Python bindings, documentation), and addresses all previous review comments. The logic flow is clean: Stats is always generated, Max is conditionally added when requested. Tensor ordering and field renaming are consistent throughout the codebase.
• No files require special attention

Important Files Changed

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart TD
    A[Fused Attention Forward Pass] --> B{generate_stats = true}
    B --> C[cuDNN SDPA Operation]
    C --> D[Always returns Stats tensor]
    D --> E{return_max_logit?}
    E -->|true| F[Also set Max output]
    E -->|false| G[Max = nullptr]
    F --> H[Output: O, Stats, Max]
    G --> I[Output: O, Stats]
    H --> J[Python Layer: aux_ctx_tensors = Stats]
    I --> K[Python Layer: aux_ctx_tensors = Stats]
    J --> L[Compute max_logit from Max tensor]
    L --> M[Return: O, aux_ctx_tensors, max_logit]
    K --> N[Return: O, aux_ctx_tensors]
    
    style D fill:#90EE90
    style F fill:#FFD700
    style H fill:#87CEEB
    style M fill:#87CEEB

_{Last reviewed commit: 56e46fd}